package com.yootoo.extpmml.utils;

import com.yootoo.extpmml.Entites.model.ModelDesc;
import com.yootoo.extpmml.Entites.ModelEvaluator;
import com.yootoo.extpmml.Exception.ModelException;
import com.yootoo.extpmml.smg3.DynamicStrategiesRunner;
import ml.dmlc.xgboost4j.java.XGBoost;
import org.dmg.pmml.PMML;
import org.drools.core.impl.InternalKnowledgeBase;
import org.drools.core.impl.KnowledgeBaseFactory;
import org.jpmml.evaluator.*;
import org.jpmml.model.PMMLUtil;
import org.kie.api.KieBase;
import org.kie.api.io.ResourceType;
import org.kie.internal.builder.KnowledgeBuilder;
import org.kie.internal.builder.KnowledgeBuilderFactory;
import org.kie.internal.io.ResourceFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.util.*;

public class ModelHelper {
    private static final Logger LOGGER = LoggerFactory.getLogger(ModelHelper.class);
    private static final ModelEvaluatorFactory modelEvaluatorFactory = ModelEvaluatorFactory.newInstance();

    public static void addToEvaluators(Map<String,ModelEvaluator> evaluators,ModelDesc modelDesc){
        try {
            ModelEvaluator evaluator = toEvaluator(modelDesc);
            evaluators.put(modelDesc.getName(),evaluator);
        }catch (Exception e){
            LOGGER.error("add to Evaluators failed",e);
        }
    }

    public static ModelEvaluator toEvaluator(ModelDesc modelDesc){
        ModelEvaluator modelEvaluator = new ModelEvaluator();
        modelEvaluator.setModel(modelDesc);
        try {
            if (modelDesc.getSource().equalsIgnoreCase("Evaluator")){
                modelEvaluator.setInnerEvaluator(modelEvaluatorFactory.newModelEvaluator(toPMML(modelDesc)));
            }else if (modelDesc.getSource().equalsIgnoreCase("Booster")){
                modelEvaluator.setInnerEvaluator(XGBoost.loadModel(toIS(modelDesc)));
            }else if (modelDesc.getSource().equalsIgnoreCase("Drools")){
                modelEvaluator.setInnerEvaluator(initKieBase(modelDesc));
            }else if (modelDesc.getSource().equalsIgnoreCase("Smg3")){
                modelEvaluator.setInnerEvaluator(new DynamicStrategiesRunner(modelDesc));
            }
        }catch (Exception e){
            throw new ModelException("模型初始化失败",e);
        }
        return modelEvaluator;
    }

    private static KieBase initKieBase(ModelDesc rule) {
        long startTime = System.currentTimeMillis();
        KnowledgeBuilder kb = KnowledgeBuilderFactory.newKnowledgeBuilder();
        kb.add(ResourceFactory.newByteArrayResource(rule.getPmml()), ResourceType.DRL);
        if (kb.hasErrors()) {
            throw new RuntimeException("Build Errors:\n" + kb.getResults().toString());
        }
        InternalKnowledgeBase kBase = KnowledgeBaseFactory.newKnowledgeBase();
        kBase.addPackages(kb.getKnowledgePackages());
        long endTime = System.currentTimeMillis();
        LOGGER.info("规则base({})初始化构建完成，耗时{}ms",rule.getName(),endTime - startTime );
        return kBase;
    }

    public static PMML toPMML(ModelDesc modelDesc){
        PMML pmml = null;
        try {
            pmml = PMMLUtil.unmarshal(toIS(modelDesc));
        } catch (Exception e) {
            throw new ModelException("非正常模型");
        }
        return pmml;
    }

    private static ByteArrayInputStream toIS(ModelDesc modelDesc){
        ByteArrayInputStream is = new ByteArrayInputStream(modelDesc.getPmml());
        return is;
    }
}
